# imports
from openmm.app import *
from openmm import *
from openmm.unit import *


from sys import stdout
import os
#os.environ["CUDA_VISIBLE_DEVICES"]="0"


# definition of methods for e-field


def create_electricforce(Ex):
    """
    This method creates force object that applies electric force onto z axis
    of the simulation. The electric force must be in the units of 
    kilojoules_per_mole/nanometer/elementary_charge.
    
    Arguments: 
        Ex as Quantity
    Returns: 
        Force object
    """
    
    Ex = Ex.in_units_of(kilojoules_per_mole/nanometer/elementary_charge)
    potential = "-q*Ex*z"
    force = CustomExternalForce(potential)
    force.addPerParticleParameter("q")
    force.addGlobalParameter("Ex", Ex)
    
    return force


def apply_electric_permolecule(system, force):
    """
    This method is used to apply electric force to a given system.
    Arguments: 
        system as System
        force as Force
    Returns:
        None
    """
    nonbonded = [f for f in system.getForces() if isinstance(f, NonbondedForce)][0]
    for i in range(system.getNumParticles()):
        charge, sigma, epsilon = nonbonded.getParticleParameters(i)
        # You have the charge for particle i, so add it to the CustomExternalForce.
        force.addParticle(i, [charge])
    system.addForce(force)
        
def voltage_per_cm_to_openmm(v_per_cm):
    """
    Coverts voltage per centimeter (as volt/centimer Quantity) to
    kilojoules_per_mole/nanometer/elementary_charge units. 
    
    Arguments: 
        v_per_cm as Quantity
    Returns
        v as Quantity 
    """
    v =  v_per_cm / item
    return v.in_units_of(kilojoules_per_mole/nanometer/elementary_charge)
        
    
    


# create NaCl system


# Creating a simulation as according
testsystem = SodiumChlorideCrystal()
system = testsystem.system
integrator = LangevinMiddleIntegrator(300*kelvin, 1/picosecond, 4*femtoseconds)
topology = testsystem.topology
positions = testsystem.positions


# apply INSANE voltage


voltage = 100000000000000000000000000000 * volt/centimeter
voltage = voltage_per_cm_to_openmm(voltage)
eforce = create_electricforce(voltage)
apply_electric_permolecule(system, eforce)


# define simulation


simulation = Simulation(topology, system, integrator)
simulation.context.setPositions(positions)
simulation.context.setVelocitiesToTemperature(300 * kelvins)



simulation.reporters.append(DCDReporter('output.dcd', 100))
simulation.reporters.append(StateDataReporter(stdout, 100, step=True,
        potentialEnergy=True, temperature=True))


#save pdb file to cur dir
from simtk.openmm.app import PDBFile, PDBxFile
filehandle = open('positions.pdb', 'w')
PDBFile.writeHeader(topology, filehandle)
PDBFile.writeModel(topology, positions, filehandle, 0)
PDBFile.writeFooter(topology, filehandle)
filehandle.close()


# start sim

simulation.step(30000)